import re
import numpy as np

# 定义日志文件路径
log_file_path = 'grpo_runable.log'

# 初始化 TensorBoard 写入器
log_dir = 'tensorboard_logs'
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter(log_dir)

# 读取日志文件内容
with open(log_file_path, 'r') as file:
    lines = file.readlines()

# 遍历日志文件的每一行
for line in lines:
    # 查找包含指定日志格式的行
    if 'iteration:' in line:
        # 提取迭代次数
        iteration_match = re.search(r'iteration: (\d+) /', line)
        if iteration_match:
            iteration = int(iteration_match.group(1))

            # 提取其他键值对
            # 匹配所有形如 "key : value" 的模式
            key_value_pairs = re.findall(r'([\w/]+) : ([\d.-]+)', line)
            print(line)
            for key, value in key_value_pairs:
                # 将值转换为浮点数
                value = float(value)
                # 写入 TensorBoard
                writer.add_scalar(key, np.array(value), iteration)
writer.flush()